use core::arch::asm;

// SV39分页相关常量
const PAGE_SIZE: usize = 4096;
const PAGE_MASK: usize = PAGE_SIZE - 1;
const ENTRY_COUNT: usize = 512; // 每级页表512个条目

// 页表项标志位
const PTE_VALID: usize = 1 << 0;     // 有效位
const PTE_READ: usize = 1 << 1;      // 可读
const PTE_WRITE: usize = 1 << 2;     // 可写
const PTE_EXEC: usize = 1 << 3;      // 可执行
const PTE_USER: usize = 1 << 4;      // 用户模式可访问
const PTE_GLOBAL: usize = 1 << 5;    // 全局映射
const PTE_ACCESSED: usize = 1 << 6;  // 已访问
const PTE_DIRTY: usize = 1 << 7;     // 已修改

// 虚拟地址位段定义 (SV39)
const VPN2_SHIFT: usize = 30;
const VPN1_SHIFT: usize = 21;
const VPN0_SHIFT: usize = 12;
const VPN2_MASK: usize = 0x1FF;  // 9位
const VPN1_MASK: usize = 0x1FF;
const VPN0_MASK: usize = 0x1FF;

// 物理地址位段定义
const PPN2_SHIFT: usize = 28;
const PPN1_SHIFT: usize = 19;
const PPN0_SHIFT: usize = 10;
const PPN2_MASK: usize = 0x3FF;  // 10位
const PPN1_MASK: usize = 0x1FF;  // 9位
const PPN0_MASK: usize = 0x1FF;  // 9位

/// 页表项结构
#[repr(transparent)]
#[derive(Clone, Copy)]
struct PageTableEntry(usize);

impl PageTableEntry {
    /// 创建新的页表项
    fn new(ppn: usize, flags: usize) -> Self {
        PageTableEntry((ppn << 10) | flags)
    }

    /// 获取页表项中的物理页号
    fn ppn(&self) -> usize {
        (self.0 >> 10) & 0x0FFF_FFFF_FFFF
    }

    /// 获取页表项标志位
    fn flags(&self) -> usize {
        self.0 & 0x3FF
    }

    /// 检查页表项是否有效
    fn is_valid(&self) -> bool {
        (self.0 & PTE_VALID) != 0
    }

    /// 设置页表项为有效
    fn set_valid(&mut self) {
        self.0 |= PTE_VALID;
    }

    /// 获取页表项指向的下一级页表地址
    fn table_address(&self) -> usize {
        (self.ppn() << 12) as usize
    }

    /// 检查是否为叶节点（包含R/W/X权限）
    fn is_leaf(&self) -> bool {
        (self.0 & (PTE_READ | PTE_WRITE | PTE_EXEC)) != 0
    }
}

/// 页表结构
#[repr(align(4096))]
struct PageTable {
    entries: [PageTableEntry; ENTRY_COUNT],
}

impl PageTable {
    /// 创建新的页表
    fn new() -> Self {
        PageTable {
            entries: [PageTableEntry(0); ENTRY_COUNT],
        }
    }

    /// 获取指定索引的页表项
    fn get_entry(&self, index: usize) -> &PageTableEntry {
        &self.entries[index]
    }

    /// 获取可变的页表项
    fn get_entry_mut(&mut self, index: usize) -> &mut PageTableEntry {
        &mut self.entries[index]
    }

    /// 根据虚拟地址获取VPN索引
    fn vpn_index(vaddr: usize, level: usize) -> usize {
        match level {
            2 => (vaddr >> VPN2_SHIFT) & VPN2_MASK,
            1 => (vaddr >> VPN1_SHIFT) & VPN1_MASK,
            0 => (vaddr >> VPN0_SHIFT) & VPN0_MASK,
            _ => panic!("Invalid page level"),
        }
    }
}

/// SV39分页管理器
struct Sv39Paging {
    root_table: PageTable,
    free_page_list: usize, // 简化的空闲页管理
}

impl Sv39Paging {
    /// 创建新的SV39分页管理器
    fn new() -> Self {
        Sv39Paging {
            root_table: PageTable::new(),
            free_page_list: 0x81000000, // 假设空闲页从这个地址开始
        }
    }

    /// 分配一页物理内存
    fn alloc_page(&mut self) -> usize {
        let page = self.free_page_list;
        self.free_page_list += PAGE_SIZE;
        page
    }

    /// 映射虚拟地址到物理地址
    fn map(&mut self, vaddr: usize, paddr: usize, flags: usize) {
        let vpn2 = PageTable::vpn_index(vaddr, 2);
        let vpn1 = PageTable::vpn_index(vaddr, 1);
        let vpn0 = PageTable::vpn_index(vaddr, 0);

        // 检查是否需要创建新的二级页表
        let need_new_l1 = !self.root_table.get_entry(vpn2).is_valid();
        if need_new_l1 {
            let new_table_addr = self.alloc_and_init_table();
            let root_entry = self.root_table.get_entry_mut(vpn2);
            *root_entry = PageTableEntry::new(new_table_addr >> 12, PTE_VALID);
        }

        // 使用 unsafe 获取二级页表的可变引用
        let table1_addr = self.root_table.get_entry(vpn2).table_address();
        let entry1 = {
            let table1 = unsafe { &mut *(table1_addr as *mut PageTable) };
            table1.get_entry_mut(vpn1)
        };

        // 检查是否需要创建新的三级页表
        let need_new_l0 = !entry1.is_valid();
        if need_new_l0 {
            let new_table_addr = self.alloc_and_init_table();
            *entry1 = PageTableEntry::new(new_table_addr >> 12, PTE_VALID);
        }

        // 使用 unsafe 获取三级页表的可变引用
        let table0_addr = entry1.table_address();
        let entry0 = {
            let table0 = unsafe { &mut *(table0_addr as *mut PageTable) };
            table0.get_entry_mut(vpn0)
        };

        // 设置叶节点页表项
        let ppn = paddr >> 12;
        *entry0 = PageTableEntry::new(ppn, flags | PTE_VALID);
    }

    /// 分配并初始化一个新的页表
    fn alloc_and_init_table(&mut self) -> usize {
        let new_table_addr = self.alloc_page();
        unsafe {
            let table_ptr = new_table_addr as *mut PageTable;
            table_ptr.write(PageTable::new());
        }
        new_table_addr
    }

    /// 启用分页
    fn enable(&self, satp: usize) {
        unsafe {
            asm!("csrw satp, {}", in(reg) satp);
            asm!("sfence.vma");
        }
    }

    /// 创建内核页表并启用分页
    fn init_paging(&mut self) {
        // 映射内核代码段 (假设内核加载在0x80200000)
        let kernel_base = 0x80200000;
        let kernel_size = 0x200000; // 2MB

        // 映射内核到虚拟地址0x80200000
        for offset in (0..kernel_size).step_by(PAGE_SIZE) {
            let vaddr = kernel_base + offset;
            let paddr = kernel_base + offset;
            self.map(vaddr, paddr, PTE_READ | PTE_WRITE | PTE_EXEC);
        }

        // 映射UART设备内存
        self.map(0x10000000, 0x10000000, PTE_READ | PTE_WRITE);

        // 设置SATP寄存器启用分页 (SV39模式)
        let root_ppn = (&self.root_table as *const PageTable as usize) >> 12;
        let satp = (8 << 60) | root_ppn; // SV39模式 (mode=8)

        // 启用分页
        self.enable(satp);
    }
}

/// 初始化SV39分页
pub fn init_sv39_paging() {
    let mut paging = Sv39Paging::new();
    paging.init_paging();
}
